import requests
import json
from typing import List, Dict, Optional, Union, Any
from datetime import datetime
import logging
import re
import os.path
from urllib.parse import quote
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

logger = logging.getLogger(__name__)
BASE_URL = "http://180.184.65.98:38880/atomgit"

# 配置请求重试机制
retry_strategy = Retry(
    total=3,
    backoff_factor=0.5,
    status_forcelist=[429, 500, 502, 503, 504],
    allowed_methods=["GET"]
)
adapter = HTTPAdapter(max_retries=retry_strategy)
session = requests.Session()
session.mount("http://", adapter)
session.mount("https://", adapter)

def _make_request(endpoint: str, params: Dict = None, timeout: int = 30) -> Union[List, Dict, None]:
    """增强的API请求处理"""
    try:
        url = f"{BASE_URL}/{endpoint}"
        response = requests.get(url, params=params, timeout=timeout)
        response.raise_for_status()
        
        if response.status_code == 200:
            return response.json()
        else:
            logger.error(f"API错误: {response.text}")
            return None
            
    except requests.exceptions.RequestException as e:
        logger.error(f"请求失败: {str(e)}")
        return None

def _extract_venue_from_filename(filename: str) -> str:
    """从文件名中提取会议名称和年份"""
    try:
        if not filename:
            return "Conference"
            
        # 标准化会议名称映射
        venue_map = {
            "conf": "Conference",
            "nips": "NeurIPS",
            "neurips": "NeurIPS",
            "icml": "ICML",
            "iclr": "ICLR",
            "aaai": "AAAI",
            "cvpr": "CVPR",
            "acl": "ACL",
            "emnlp": "EMNLP",
            "naacl": "NAACL",
            "eccv": "ECCV",
            "ijcai": "IJCAI",
            "kdd": "KDD",
            "www": "WWW",
            "sigir": "SIGIR",
            "arxiv": "arXiv preprint",
            "preprint": "Preprint"
        }
        
        # 从文件名中提取会议名称和年份
        basename = os.path.basename(filename).lower()
        venue = "Conference"
        
        # 提取会议名称
        for key, value in venue_map.items():
            if key in basename:
                venue = value
                break
        
        # 提取年份
        year_match = re.search(r'(20\d{2})', basename)
        if year_match:
            year = year_match.group(1)
            return f"{venue} {year}"
        
        return venue
        
    except Exception as e:
        logger.error(f"Venue extraction failed: {str(e)}")
        return "Conference"

def _extract_year(text: str) -> int:
    """增强的年份提取"""
    current_year = datetime.now().year
    try:
        year_match = re.search(r'(20\d{2})', text)
        if year_match and 1990 <= int(year_match.group(1)) <= current_year:
            return int(year_match.group(1))
        return current_year
    except Exception as e:
        logger.error(f"年份提取失败: {str(e)}")
        return current_year

def search_papers(query: str, top_k: int = 30) -> List[Dict]:
    """根据文本查询搜索论文片段
    
    Args:
        query: 自然语言查询语句
        top_k: 返回结果数量（范围：1-1000）
        
    Returns:
        包含论文片段的列表，每个片段包含distance和entity信息
    """
    return _make_request("search_papers", {
        "query": query,
        "top_k": min(max(1, top_k), 1000)  # 确保在有效范围内
    }) or []

def query_by_paper_id(paper_id: str, top_k: int = 5) -> List[Dict]:
    """根据论文ID查询片段"""
    return _make_request("query_by_paper_id", {
        "paper_id": paper_id,
        "top_k": top_k
    }) or []

def query_by_title(title: str, top_k: int = 100) -> List[Dict]:
    """根据论文精确标题查询片段
    
    Args:
        title: 论文完整标题（需要100%精确匹配）
        top_k: 返回结果数量
        
    Returns:
        包含论文片段的列表
    """
    return _make_request("query_by_title", {
        "title": title,
        "top_k": top_k
    }) or []

def query_by_title_contain(title: str, top_k: int = 100) -> List[Dict]:
    """根据论文标题关键词查询片段（模糊匹配）
    
    Args:
        title: 标题关键词（支持模糊匹配）
        top_k: 返回结果数量
        
    Returns:
        包含论文片段的列表
    """
    return _make_request("query_by_title_contain", {
        "title": title,
        "top_k": top_k
    }) or []

def get_metadata() -> Dict:
    """获取元数据信息"""
    return _make_request("metadata") or {}

def query_whole_paper(title: Optional[str] = None, paper_id: Optional[str] = None) -> Optional[Dict]:
    """统一的论文全文查询接口
    
    Args:
        title: 论文完整标题（与paper_id二选一）
        paper_id: 论文唯一标识符（与title二选一）
        
    Returns:
        包含论文全文和元数据的字典
    """
    if paper_id:
        return _make_request("query_whole_text_by_id", {"paper_id": paper_id})
    elif title:
        return _make_request("query_whole_text_by_title", {"title": title})
    return None

def query_keywords(title: Optional[str] = None, paper_id: Optional[str] = None) -> Optional[Dict]:
    """统一的关键词查询接口
    
    Args:
        title: 论文完整标题（与paper_id二选一）
        paper_id: 论文唯一标识符（与title二选一）
        
    Returns:
        包含关键词列表和置信度分数的字典
    """
    if paper_id:
        return _make_request("query_keywords_by_id", {"paper_id": paper_id})
    elif title:
        return _make_request("query_keywords_by_title", {"title": title})
    return None

def get_keywords_stats() -> Dict[str, int]:
    """获取关键词统计信息
    
    Returns:
        关键词及其关联论文数量的字典
    """
    return _make_request("keywords_metadata") or {}

def query_by_chunk_contain(chunk: str, top_k: int = 100) -> List[Dict]:
    """在论文全文范围内检索包含特定内容的片段
    
    Args:
        chunk: 检索内容（支持布尔检索语法：AND/OR/NOT和通配符）
        top_k: 返回结果数量
        
    Returns:
        包含匹配片段的列表
    """
    return _make_request("query_by_chunk_contain", {
        "chunk": chunk,
        "top_k": top_k
    }) or []

def query_paper_metadata_that_title_contain(title: str, top_k: int = 100) -> List[List[str]]:
    """获取标题包含特定关键词的论文基础信息
    
    Args:
        title: 标题关键词（支持模糊匹配）
        top_k: 最大返回结果数量
        
    Returns:
        论文元组列表，每个元组包含[论文ID, 完整标题]
    """
    return _make_request("query_paper_metadata_that_title_contain", {
        "title": title,
        "top_k": top_k
    }) or []

def titles_like(title: str, top_k: int = 50) -> List[str]:
    """获取相似标题列表
    
    Args:
        title: 查询标题
        top_k: 返回结果数量
        
    Returns:
        相似标题列表
    """
    return _make_request("titles_like", {
        "title": title,
        "top_k": top_k
    }) or []

def query_by_keyword(keyword: str, top_k: int = 100) -> List[List[str]]:
    """根据关键词获取论文ID和标题列表
    
    Args:
        keyword: 关键词
        top_k: 最大返回结果数量
        
    Returns:
        论文元组列表，每个元组包含[论文ID, 完整标题]
    """
    return _make_request("query_by_keyword", {
        "keyword": keyword,
        "top_k": top_k
    }) or []